-- file: ray-tracer.hs

import System.Environment

sq :: Double -> Double
sq x = x * x

mag :: Double -> Double -> Double -> Double
mag x y z = sqrt (sq x + sq y + sq z)

unit_vector :: Double -> Double -> Double -> (Double, Double, Double)
unit_vector x y z = (x/d , y/d, z/d)
    where d = mag x y z

data Point = Point Double Double Double
    deriving (Eq, Show)

data Direction = Direction Double Double Double
    deriving (Eq, Show)

distance :: Point -> Point -> Double
distance (Point x1 y1 z1) (Point x2 y2 z2) = mag (x1-x2) (y1-y2) (z1-z2)

minroot :: Double -> Double -> Double -> Maybe Double
minroot a b c
    | zerop a           = Just ((-c) / b)
    | not (minusp disc) = Just (min (((-b)+discrt)/(2*a)) (((-b)-discrt)/(2*a)))
    | otherwise         = Nothing
    where disc     = (sq b) - (4 * a * c)
          discrt   = sqrt disc
          zerop x  = abs x < eps
          minusp x = x < 0.0
          eps      = 0.0001

--data Surface = Surface Color

data Color = Color Double
    deriving (Show)

eye :: Point
eye = Point 0 0 200

tracer :: FilePath -> [Surface] -> Light -> Int -> IO()
tracer pathname world light res =
    let
        cols = [color_at world light (pos i) (pos j)  | j <- [0..(res*100)-1],
                                                        i <- [0..(res*100)-1]]
    in
       do putPGM pathname (res*100) (res*100) cols
    where 
          putPGM :: FilePath -> Int -> Int -> [Int] -> IO()
          putPGM pathname w h cols = 
              let header = "P2 " ++ show w ++ " " ++ show h ++ " 255\n"
                  body   = foldr (++) "" [(show c) ++ "\n" | c <- cols]
              in do writeFile pathname (header ++ body)
          
          pos :: Int -> Double
          pos i = (fromIntegral i) / (fromIntegral res) - 50.0

color_at :: [Surface] -> Light -> Double -> Double -> Int
color_at world light x y =
    let
        Point ex ey ez = eye
        (xr, yr, zr)   = unit_vector (x-ex) (y-ey) (0-ez)
    in
        color (sendray world light eye xr yr zr)
    where
        color (Color col) = truncate (col * 255)

sendray :: [Surface] -> Light -> Point ->
           Double -> Double -> Double ->
           Color
sendray world light pt xr yr zr =
    case first_hit world pt xr yr zr of
      Just (s@(Sphere (Color col) _ _ _), int)
        -> case is_specular s of
             True  -> let Direction x y z = reflection_vector s int (Direction xr yr zr)
                      in  sendray world light int x y z
             False -> case shadowed int of
                        True  -> Color 0.0
                        False -> Color ((lambert s int light) * col)
      Nothing
        -> Color 0.0
    where shadowed :: Point -> Bool
          shadowed int@(Point xi yi zi) =
              let PointLight (Point xl yl zl) = light
                  (xr', yr', zr') = unit_vector (xi-xl) (yi-yl) (zi-zl)
                  eps = 0.0001
              in case first_hit world (Point xl yl zl) xr' yr' zr' of
                   Just (_, light_int) -> if distance int light_int < eps
                                          then False
                                          else True

reflection_vector :: Surface -> Point -> Direction -> Direction
reflection_vector s p d = let Point px py pz     = p
                              Direction dx dy dz = d
                              a                  = px*dx + py*dy + pz*dz
                              (nx,ny,nz)         = normal s p
                              (dx',dy',dz')      = unit_vector px-2*a*nx py-2*a*ny pz-2*a*nz
                          in Direction dx' dy' dz'

first_hit :: [Surface] -> Point ->
             Double -> Double -> Double ->
             Maybe (Surface, Point)
first_hit world pt xr yr zr = nearest (map hit world)
    where nearest :: [Maybe (Double, Surface, Point)] -> Maybe (Surface, Point)
          nearest hits = case foldr cmp Nothing hits of
                           Just (d, s, h) -> Just (s, h)
                           otherwise      -> Nothing
          
          cmp :: Maybe (Double, Surface, Point) ->
                 Maybe (Double, Surface, Point) ->
                 Maybe (Double, Surface, Point)
          cmp Nothing Nothing = Nothing
          cmp x       Nothing = x
          cmp Nothing y       = y
          cmp x@(Just (xd, _, _)) y@(Just (yd, _, _))
            | xd < yd   = x
            | otherwise = y
          
          hit :: Surface -> Maybe (Double, Surface, Point)
          hit s = case intersect s pt xr yr zr of 
                    Just h  -> Just (distance h pt, s, h)
                    Nothing -> Nothing

lambert :: Surface -> Point -> Light -> Double
lambert s int@(Point xp yp zp) (PointLight (Point xl yl zl)) =
    let (xn, yn, zn) = normal s int
        (xr, yr, zr) = unit_vector (xp-xl) (yp-yl) (zp-zl)
    in max 0 (xr*xn + yr*yn + zr*zn)

data Surface = Sphere {
      color    :: Color
    , radius   :: Double
    , center   :: Point
    , specular :: Bool
    } deriving (Show)

data Light   = PointLight Point
    deriving (Show)

is_specular :: Surface -> Bool
is_specular (Sphere _ _ _ s) = s

intersect :: Surface -> Point -> Double -> Double -> Double -> Maybe Point
intersect (Sphere _ r (Point cx cy cz) _) (Point px py pz) xr yr zr =
    let a = (sq xr) + (sq yr) + (sq zr)
        b = 2 * ((px-cx)*xr + (py-cy)*yr + (pz-cz)*zr)
        c = (sq (px-cx)) + (sq (py-cy)) + (sq (pz-cz)) - (sq r)
    in case minroot a b c of
         Just n  -> Just (Point (px+xr*n) (py+yr*n) (pz+zr*n))
         Nothing -> Nothing

normal :: Surface -> Point -> (Double, Double, Double)
normal (Sphere _ _ (Point cx cy cz) _) (Point px py pz) =
    unit_vector (cx-px) (cy-py) (cz-pz)

ray_test :: Int -> IO()
ray_test res =
    let world = [Sphere (Color 0.8) 200 (Point 0 200 (-1200)) True,
                 Sphere (Color 0.7) 100 (Point 0 (-200) (-1100)) False]
        light = PointLight (Point 1200 0 (-1200))
    in do tracer "spheres.pgm" world light res

main = do args <- getArgs
          case args of
            []   -> ray_test 3
            x:[] -> ray_test (read x)
